"""
任务 房价价格预测

"""
import pandas as pd
from my_utils import *
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression

DATA_FILE = ".\data\house_data.csv"

# 使用的特征列
"""
bedrooms:房间个数  bathrooms:卫生间个数 sqft_living:居住平米数 sqft_lot:总平米数
sqft_above: 上层平米数 sqft_basement: 底层平米数
"""
FEAT_COLS = ['bedrooms','bathrooms','sqft_living','sqft_lot','sqft_above','sqft_basement']

def main():
    house_data = pd.read_csv(DATA_FILE, usecols=FEAT_COLS + ['price'])
    ai_utils.plot_feat_and_price(house_data)
    X = house_data[FEAT_COLS].values  # 二维数组
    y = house_data['price'].values  # 一维数组
    # test_size 和 random_state 是什么意思
    # 分割数据集
    X_train, X_test, y_train, y_test = train_test_split(X,y,test_size=1/3, random_state=10)
    # print("X_train",X_train,type(X_train))
    # print("X_test", X_test, type(X_test))
    # print("y_train", y_train, type(y_train))
    # print("y_test", y_test, type(y_test))
    # 建立线性回归模型
    linear_reg_model = LinearRegression()
    # 模型训练
    linear_reg_model.fit(X_train,y_train)
    # 验证模型
    r2_score = linear_reg_model.score(X_test,y_test)
    print('模型的R2值', r2_score) # 0.5058817864948637

    #单个样本房价预测
    i = 50
    single_test_feat = X_test[i, :]
    y_true = y_test[i]
    y_pred = linear_reg_model.predict([single_test_feat])
    print("样本特征:", single_test_feat)
    print('真实价格: {}, 预测价格： {}'.format(y_true,y_pred))

if __name__ =='__main__':
    main()